import re
import pickle

def step2(file_path):
    result_sent, sent, record, idx = [], [], [], []
    count = 0
    with open(file_path,'r',encoding='gbk',errors='ignore') as file:
        for line in file:
            ret = re.search('(\<\d+\>)(.+)',line.strip())
            if ret:
                count = ret.group(1)
                sentence = ret.group(2)
                pos = sentence.split(' ')
                fenci = re.sub('\/[a-zA-Z]*','',sentence).split(' ')
                sent.append(''.join(i for i in fenci))
                num,i=[],0   
                for k in fenci:
                    num.append(list(range(i,i+len(k))))
                    i=i+len(k) 
            elif line.strip() == "====PGraph Beg====":
                continue
            elif line.strip() == "====PGraph End====":
                sent.append(num)
                sent.append(idx)
                sent.append(record)
                sent.append(pos)
                result_sent.append(sent)
                sent, record, idx = [], [], []
            elif len(line.strip()) != 0:
                head = line.strip().split(',')
                # if head[-1] in ['sbj_HX', 'obj_HX']:
                record.append((head[0],head[-1],head[2]))
                x = head[1].strip('()').split(':')
                y = head[3].strip('()').split(':')
                if len(set(x)) == 1 :
                    h = [int(i) for i in x[1:]]
                else:
                    m, n = int(x[0]), int(x[1])
                    h = list(range(m, m+(n-m)+1))
                if len(set(y)) == 1 :
                    t = [int(i) for i in y[1:]]
                else:
                    m, n = int(y[0]), int(y[1])
                    t = list(range(m, m+(n-m)+1))
                idx.append((num.index(h),head[-1],num.index(t)))
    return result_sent

if __name__=="__main__":
    train = 'final/'+ 'train_data.pkl'
    val = 'final/'+ 'val_data.pkl'
    test = 'final/'+ 'test_data.pkl'
    path = "50w_all_2.out"
    result_sent = step2(path)
    print(result_sent[:2])
    print(len(result_sent))
    with open(train,'wb') as out_file:
        pickle.dump(result_sent[:400000],out_file) 
    with open(val,'wb') as out_file:
        pickle.dump(result_sent[400000:401000],out_file) 
    with open(test,'wb') as out_file:
        pickle.dump(result_sent[401000:402000],out_file) 

    multi_v, single_v = [], []
    for item in result_sent[402000:408000]:
        v = set()
        for i in item[2]:
            v.add(i[0])
        if len(v) == 1:
            single_v.append(item)
        else:
            multi_v.append(item)
    print("len(single_v):", len(single_v))
    print("len(multi_v):", len(multi_v))
    print(single_v[:2])
    with open('final/single.pkl','wb') as out_file:
        pickle.dump(single_v[:1000],out_file)
    m = multi_v[:1000]
    m.append(result_sent[66926])
    print(len(m))
    print(m[-1])
    print(result_sent[66926])
    with open('final/multi.pkl','wb') as out_file:
        pickle.dump(m,out_file)
    